#!/usr/bin/env python

import matplotlib.pyplot as plt
import math

from sympy import diff
from sympy.abc import x
from sympy.core import Symbol
from sympy.solvers import solve

from choose_functions_number import choose_functions_number

xmin = -1.0
xmax = 1.0

a = math.sin(3)
b = math.cos(3)


def colocations(steps):
   basis = [(x ** (2*i-2) * (1 - x ** 2)) for i in range(0, steps + 1)]
   basis[0] = 0

   xs = [xmin + (0 - xmin) / (steps + 1) * i for i in range(1, steps + 1)]
   ax = [1] + [Symbol('a%i' % i) for i in range(1, steps + 1)]

   yapprox = sum(ax[i] * basis[i] for i in range(0, steps + 1))
   error = a * diff(diff(yapprox, 'x'), 'x') + (1 + b * x ** 2) * yapprox + 1

   matrix = [error.subs(x, xs[i]) for i in range(0, steps)]

   sol = solve(matrix)
   y = yapprox.subs(sol)

   print 'Colocations method. Interval [{0},{1}] with {2} points:'.format(xmin, xmax, steps)
   for xi in xs:
      print '{0:+5.3f}\t: {1}'.format(xi, y.subs(x,xi))
   print

   print y
   return y


def plot_colocations(**kwargs):
    n = choose_functions_number(colocations, 0.001)
    y = colocations(n)
    gsteps = 100
    xs = [xmin + (xmax - xmin) / gsteps * i for i in range(0, gsteps + 1)]
    ys = [y.subs(x, xi) for xi in xs]

    plt.plot(xs, ys, **kwargs)

plot_colocations(color='red')
plt.show()
